use std::{
    collections::HashMap, sync::{atomic::{AtomicU64, Ordering}, Arc}, time::Duration
};

use anyhow::Error;
use bincode::{Decode, Encode};
use std::time::SystemTime;
use crate::{args::Args, command::Command, frame::Frame, persistence::rdb_file::RdbFile};
use tokio::sync::{
    mpsc::{channel, Receiver, Sender},
    oneshot,
};

#[derive(Clone, Encode, Decode)]
pub struct DatabaseSnapshot {
    pub collections: HashMap<String, VectorCollection>,
}

/**
 * 消息
 *
 * @param sender 发送者
 * @param command 命令
 */
pub enum DatabaseMessage {
    Changes(oneshot::Sender<u64>),
    Command{sender: oneshot::Sender<Frame>, command: Command},
    Snapshot(oneshot::Sender<DatabaseSnapshot>),
    Restore(DatabaseSnapshot)
}

impl Default for DatabaseSnapshot {
    fn default() -> Self {
        Self {
            collections: HashMap::with_capacity(100000),
        }
    }
}

/**
 * DB 管理器
 */
pub struct DatabaseManager {
    senders: Vec<Sender<DatabaseMessage>>
}

impl DatabaseManager {

    /**
     * 创建 DB 管理器
     *
     * @param config 参数
     */
    pub fn new(args: Arc<Args>) -> Self {

        let mut dbs = Vec::new();
        let mut senders = Vec::new();
        let mut rdb_file = RdbFile::new(args.dbfilename.clone());
        let _ = rdb_file.load();

        for id in 0..args.databases {
            let db = Db::new(rdb_file.get_database(id));
            senders.push(db.sender.clone());
            dbs.push(db);
        }

        for mut db in dbs {
            tokio::spawn(async move {
                db.run().await;
            });
        }

        let args_clone = args.clone();
        let senders_clone = senders.clone();

        tokio::spawn(async move {
            let period = Duration::from_secs_f64(1.0 / args_clone.hz);
            let mut interval = tokio::time::interval(period);
            loop {

                interval.tick().await;

                let mut changes = 0;
                for sender in &senders_clone {
                    let (tx, rx) = oneshot::channel();
                    if sender.send(DatabaseMessage::Changes(tx)).await.is_ok() {
                        if let Ok(count) = rx.await {
                            changes += count;
                        }
                    }
                }

                let should_save = {
                    let now = SystemTime::now();
                    let elapsed = now.duration_since(rdb_file.last_save_time).unwrap().as_secs();
                    args_clone.save.iter().any(|rule| {
                        elapsed >= rule.seconds && changes.saturating_sub(rdb_file.last_save_changes) >= rule.changes
                    })
                };

                if should_save {
                    let mut snapshots = Vec::new();
                    for sender in &senders_clone {
                        let (tx, rx) = oneshot::channel();
                        if sender.send(DatabaseMessage::Snapshot(tx)).await.is_ok() {
                            if let Ok(snapshot) = rx.await {
                                snapshots.push(snapshot);
                            }
                        }
                    }

                    for (index, snapshot) in snapshots.into_iter().enumerate() {
                        rdb_file.set_database(index, snapshot); // 更新快照
                    }
                    rdb_file.last_save_changes = changes;
                    let _ = rdb_file.save();
                }
            }
        });
        DatabaseManager { 
            senders 
        }
    }

    /**
     * 获取发送者
     *
     * @param idx 数据库索引
     */
    pub fn get_sender(&self, idx: usize) -> Sender<DatabaseMessage> {
        if let Some(sender) = self.senders.get(idx) {
            sender.clone()
        } else {
            panic!("Index out of bounds");
        }
    }

    /**
     * 获取发送者【所有】
     */
    pub fn get_senders(&self) -> Vec<Sender<DatabaseMessage>> {
        self.senders.clone()
    }
}

#[derive(Clone, Encode, Decode)]
pub struct Vector{
    pub vector: Vec<f32>, // 向量
    pub metadata: HashMap<String, String>, // 元数据
    pub norm: f32, // 范数
}

#[derive(Clone, Encode, Decode)]
pub struct VectorCollection {
    pub dimension: usize,
    pub vectors: HashMap<String, Vector>,
}

impl VectorCollection {
    pub fn new(dimension: usize) -> Self {
        VectorCollection {
            dimension,
            vectors: HashMap::new(),
        }
    }

    pub fn contains_id(&self, id: &str) -> bool {
        self.vectors.contains_key(id)
    }
}

/**
 * 数据库
 * 
 * @param receiver
 * @param sender 发送者
 * @param collections 向量集
 * @param changes 
 */
pub struct Db {
    receiver: Receiver<DatabaseMessage>,
    sender: Sender<DatabaseMessage>,
    pub collections: HashMap<String, VectorCollection>,
    changes: AtomicU64,
}

impl Db {

    /**
     * 创建数据库
     * 
     * @param args 数据库配置
     */
    pub fn new(snapshot: DatabaseSnapshot) -> Self {

        let (sender, receiver) = channel(1024);
        let collections = snapshot.collections;

        Db {
            collections,
            changes: AtomicU64::new(0),
            receiver,
            sender,
        }
    }

    /**
     * 运行数据库
     * 
     * 1. 命令执行
     * 
     * @param self 本身
     */
    async fn run(&mut self) {
        loop {
            match self.receiver.recv().await {
                Some(DatabaseMessage::Command { sender, command }) => {
                    let result: Result<Frame, Error> = match command {
                        Command::CreateCollection(create_collection) => create_collection.apply(self),
                        Command::DeleteCollection(delete_collection) => delete_collection.apply(self),
                        Command::Search(search) => search.apply(self),
                        Command::Insert(insert) => insert.apply(self),
                        Command::Delete(delete) => delete.apply(self),
                        _ => Err(Error::msg("Unknown command")),
                    };

                    match result {
                        Ok(f) => {
                            let _ = sender.send(f);
                        },
                        Err(e) => eprintln!("Error applying command: {:?}", e),
                    }
                },
                Some(DatabaseMessage::Changes(sender)) => {
                    let count = self.changes.load(Ordering::Relaxed);
                    let _ = sender.send(count);
                },
                Some(DatabaseMessage::Restore(snapshot)) => {
                    self.collections = snapshot.collections;
                },
                Some(DatabaseMessage::Snapshot(sender)) => {
                    let snapshot = DatabaseSnapshot {
                        collections: self.collections.clone()
                    };
                    let _ = sender.send(snapshot);
                },
                None => {}
            }
        }
    }

    /**
     * 保存键值
     *
     * @param key 键名
     * @param value 值
     */
    pub fn create_collection(&mut self, key: String, value: VectorCollection) {
        self.collections.insert(key, value);
    }

    /**
     * 删除键值
     *
     * @param key 键名
     * @return 如果删除成功，返回被删除的值；如果删除失败，返回 None
     */
    pub fn remove_collection(&mut self, key: &str) -> Option<VectorCollection> {
        if self.collections.contains_key(key) {
           self.collections.remove(key)
        } else {
           None
        }
    }

    /**
     * 获取键值
     *
     * @param key 键名
     */
    pub fn get(&mut self, key: &str) -> Option<&VectorCollection> {
        self.collections.get(key)
    }

    /**
     * 获取键值【引用】
     *
     * @param key 键名
     */
    pub fn get_mut(&mut self, key: &str) -> Option<&mut VectorCollection> {
        self.collections.get_mut(key)
    }
}